"""Bandwidth management - speed limiting and traffic control."""

import asyncio
import time
from dataclasses import dataclass

import psutil
from loguru import logger


@dataclass
class BandwidthStats:
    """Bandwidth statistics."""

    download_speed: float  # bytes/s
    upload_speed: float  # bytes/s
    total_downloaded: int  # bytes
    total_uploaded: int  # bytes
    active_connections: int


class BandwidthManager:
    """Manages global bandwidth for speed limiting and traffic control."""

    def __init__(
            self,
            max_download_speed: int | None = None,  # bytes/s
            max_upload_speed: int | None = None,  # bytes/s
            enable_idle_mode: bool = False,
            idle_threshold_minutes: int = 10,
    ):
        """Initializes the BandwidthManager.

        Args:
            max_download_speed (int | None): Maximum download speed in bytes/s.
            max_upload_speed (int | None): Maximum upload speed in bytes/s.
            enable_idle_mode (bool): Whether to enable idle mode.
            idle_threshold_minutes (int): Idle threshold in minutes.
        """
        self.max_download_speed = max_download_speed
        self.max_upload_speed = max_upload_speed
        self.enable_idle_mode = enable_idle_mode
        self.idle_threshold_minutes = idle_threshold_minutes

        # Statistics
        self._total_downloaded = 0
        self._total_uploaded = 0
        self._active_connections = 0

        # Rate limiter
        self._download_tokens = 0.0
        self._upload_tokens = 0.0
        self._last_refill = time.monotonic()

        # System monitoring
        self._last_idle_check = time.time()
        self._is_idle = False

        logger.info(
            f"Bandwidth manager initialized - "
            f"Download limit: {self._format_speed(max_download_speed)}, "
            f"Upload limit: {self._format_speed(max_upload_speed)}"
        )

    async def acquire_download_permit(self, bytes_count: int) -> None:
        """Acquires a download permit (rate limiting).

        Args:
            bytes_count (int): The number of bytes to download.
        """
        if self.max_download_speed is None:
            return  # Unlimited

        # Check if the system is idle
        if self.enable_idle_mode:
            await self._check_system_idle()
            if self._is_idle:
                return  # Unlimited when idle

        # Token bucket algorithm
        await self._wait_for_tokens(bytes_count, is_download=True)
        self._total_downloaded += bytes_count

    async def acquire_upload_permit(self, bytes_count: int) -> None:
        """Acquires an upload permit (rate limiting).

        Args:
            bytes_count (int): The number of bytes to upload.
        """
        if self.max_upload_speed is None:
            return  # 无限制

        if self.enable_idle_mode and self._is_idle:
            return

        await self._wait_for_tokens(bytes_count, is_download=False)
        self._total_uploaded += bytes_count

    async def _wait_for_tokens(self, bytes_count: int, is_download: bool) -> None:
        """Waits for the token bucket to refill.

        Args:
            bytes_count (int): The number of bytes needed.
            is_download (bool): True if for download, False if for upload.
        """
        max_speed = self.max_download_speed if is_download else self.max_upload_speed
        if max_speed is None:
            return

        while True:
            # Refill tokens
            now = time.monotonic()
            elapsed = now - self._last_refill
            self._last_refill = now

            tokens_to_add = max_speed * elapsed

            if is_download:
                self._download_tokens = min(
                    self._download_tokens + tokens_to_add, max_speed * 2
                )  # Store up to 2 seconds of tokens
                if self._download_tokens >= bytes_count:
                    self._download_tokens -= bytes_count
                    break
            else:
                self._upload_tokens = min(self._upload_tokens + tokens_to_add, max_speed * 2)
                if self._upload_tokens >= bytes_count:
                    self._upload_tokens -= bytes_count
                    break

            # Wait for a short period before retrying
            await asyncio.sleep(0.01)

    async def _check_system_idle(self) -> None:
        """Checks if the system is idle."""
        now = time.time()

        # Check every 30 seconds
        if now - self._last_idle_check < 30:
            return

        self._last_idle_check = now

        try:
            # Check CPU usage
            cpu_percent = psutil.cpu_percent(interval=1)

            # Check disk IO
            disk_io = psutil.disk_io_counters()

            # Check network IO
            net_io = psutil.net_io_counters()

            # Simple idle judgment: CPU < 20% and no large IO recently
            self._is_idle = cpu_percent < 20

            logger.debug(f"System idle check: CPU={cpu_percent:.1f}%, Idle={self._is_idle}")

        except Exception as e:
            logger.error(f"System idle check failed: {e}")
            self._is_idle = False

    def register_connection(self) -> None:
        """Registers an active connection."""
        self._active_connections += 1

    def unregister_connection(self) -> None:
        """Unregisters an active connection."""
        self._active_connections = max(0, self._active_connections - 1)

    def get_stats(self) -> BandwidthStats:
        """Gets current bandwidth statistics.

        Returns:
            BandwidthStats: Current bandwidth statistics.
        """
        # Calculate current speed (simplified, should use a sliding window in practice)
        return BandwidthStats(
            download_speed=0.0,  # TODO: Implement actual speed calculation
            upload_speed=0.0,
            total_downloaded=self._total_downloaded,
            total_uploaded=self._total_uploaded,
            active_connections=self._active_connections,
        )

    def set_download_limit(self, speed_bytes_per_sec: int | None) -> None:
        """Sets the maximum download speed.

        Args:
            speed_bytes_per_sec (int | None): The maximum download speed in bytes per second, or None for unlimited.
        """
        self.max_download_speed = speed_bytes_per_sec
        logger.info(f"Download limit updated: {self._format_speed(speed_bytes_per_sec)}")

    def set_upload_limit(self, speed_bytes_per_sec: int | None) -> None:
        """Sets the maximum upload speed limit for the manager.

        Args:
            speed_bytes_per_sec (int | None): The maximum upload speed in bytes per second, or None for unlimited.
        """
        self.max_upload_speed = speed_bytes_per_sec
        logger.info(f"Upload limit updated: {self._format_speed(speed_bytes_per_sec)}")

    @staticmethod
    def _format_speed(bytes_per_sec: int | None) -> str:
        """Formats the speed for display.

        Args:
            bytes_per_sec (int | None): Speed in bytes per second.

        Returns:
            str: Formatted speed string (e.g., "10.24 MB/s").
        """
        if bytes_per_sec is None:
            return "Unlimited"

        units = [("B/s", 1), ("KB/s", 1024), ("MB/s", 1024 ** 2), ("GB/s", 1024 ** 3)]

        for unit, divisor in reversed(units):
            if bytes_per_sec >= divisor:
                return f"{bytes_per_sec / divisor:.2f} {unit}"

        return f"{bytes_per_sec} B/s"


class SpeedLimiter:
    """A simple speed limiter using a token bucket algorithm."""

    def __init__(self, max_speed: Optional[int] = None):
        """Initializes the SpeedLimiter.

        Args:
            max_speed (Optional[int]): Maximum speed in bytes per second, or None for unlimited.
        """
        self.max_speed = max_speed
        self._tokens = 0.0
        self._last_refill = time.monotonic()

    async def consume(self, bytes_count: int) -> None:
        """Consumes a specified number of bytes, blocking until enough tokens are available.

        Args:
            bytes_count (int): The number of bytes to consume.
        """
        if self.max_speed is None:
            return  # Unlimited

        while True:
            # Refill tokens
            now = time.monotonic()
            elapsed = now - self._last_refill
            self._last_refill = now

            tokens_to_add = self.max_speed * elapsed
            self._tokens = min(self._tokens + tokens_to_add, self.max_speed * 2)

            # Consume tokens
            if self._tokens >= bytes_count:
                self._tokens -= bytes_count
                break

            # Wait
            await asyncio.sleep(0.01)

    def set_speed(self, speed: int | None) -> None:
        """Sets the speed limit.

        Args:
            speed (int | None): The new speed limit in bytes per second, or None for unlimited.
        """
        self.max_speed = speed
        self._tokens = 0.0
        self._last_refill = time.monotonic()


# 全局带宽管理器实例
_global_bandwidth_manager: BandwidthManager | None = None


def get_bandwidth_manager() -> BandwidthManager:
    """Retrieves the global BandwidthManager instance.

    Returns:
        BandwidthManager: The global BandwidthManager instance.
    """
    global _global_bandwidth_manager
    if _global_bandwidth_manager is None:
        _global_bandwidth_manager = BandwidthManager()
    return _global_bandwidth_manager


def set_bandwidth_manager(manager: BandwidthManager) -> None:
    """Sets the global BandwidthManager instance.

    Args:
        manager: The BandwidthManager instance to set.
    """
    global _global_bandwidth_manager
    _global_bandwidth_manager = manager
